# -*- encoding: utf-8 -*-
# Python 3.6.7rc1 64bit
'''
@File        : PdfDirGenerator.py
@Time        : 2022/02/02 14:50
@Author      : Wreng
@Description : 根据包含目录信息的 txt 来自动给pdf添加书签
@Other       : version - 3.6.7rc1 64bit,  PyPDF2 1.26.0
@Note        :
                * PyPDF2 is required. You can install by the command 'pip install pypdf2'.
                * You must run the code with a version of Python before v3.7.
                * If you encounter the error: ValueError: {'/Type': '/Outlines', '/Count': 0} is not in list,
                  please modify the function <getOutlineRoot> in ${PYTHON_PATH}/site-packages/PyPDF2/pdf.py).
                  You could refer more from https://www.codetd.com/en/article/11823498.
                * You can run the program by typing following code on the command line:
                
                    python <path to pdf> <path to txt> <offset> [output path] [level mark]

                  e.g. python .\PdfDirGenerator.py .\多源多目标统计信息融合Mahler.pdf .\dir.txt 27
                    - The first parameter is the path to pdf to mark.
                    - The second parameter is the path to txt containing directory information.
                      It's a pity that you might do it manually or use OCR tools.
                    - The third parameter is the number of offset page to the text.
                    - The fourth optional parameter is the path to the output.
                    - The fifth optional parameter is the marker. The default value is '.'.
'''

import PyPDF2
import sys

class PdfDirGenerator:

    def __init__(self, pdf_path:str, txt_path:str, offset:int, out_path:str=None, levelmark:str='.'):
        
        self.pdf_path = pdf_path    # pdf路径
        self.txt_path = txt_path    # 包含pdf目录信息的txt
        self.offset = offset        # 目录页数偏移量
        self.out_path = out_path    # 输出路径
        self.levelmark = levelmark  # 用于判断书签级别的标志符

        # 由于书签是有序的，我们只需保存/更新离子书签最近的父节点
        self.dir_parent = [None]   

    def getLevelId(self, level):
        """计算书签的级数（级数的标志符号为“.”）
        一级目录: 0 个“.”，例如: 第1章、附录A等
            二级目录: 1个“.”，例如: 1.1、A.1
                三级目录: 2个“.”，例如: 2.1.3
        """
        mark_num = 0
        for c in level:
            if c == self.levelmark:
                mark_num += 1
        return mark_num + 1

    def run(self):
        
        print("--------------------------- Adding the bookmark ---------------------------")
        print(" * PDF Source: %s" % self.pdf_path)
        print(" * TXT Source: %s" % self.txt_path)
        print(" * Offset: %d" % self.offset)
        print("---------------------------------------------------------------------------")
        with open(self.txt_path, 'r', encoding='utf-8') as txt:
            
            pdf_reader = PyPDF2.PdfFileReader(self.pdf_path)
            pdf_writer = PyPDF2.PdfFileWriter()
            
            pdf_writer.cloneDocumentFromReader(pdf_reader)
            # BUG: ValueError: {’/Type’: ‘/Outlines’, ‘/Count’: 0} is not in list
            # 修改代码 ${PYTHON_PATH}/site-packages/PyPDF2/pdf.py): getOutlineRoot 函数
            # 参考：https://www.codetd.com/en/article/11823498

            lines = txt.readlines()
            num_all_lines = len(lines)
            for i, line in enumerate(lines):
                # pline = line.split(' ')
                pline = line.split(None, maxsplit=-1)
                level = pline[0]; title = pline[1]; page = int(pline[2]) + self.offset

                # 1. 计算当前的 level 的级数 id
                # 2. 插入书签的父结点存放在 dir_parent[id-1] 上
                # 3. 更新 dir_parent[id] 
                id = self.getLevelId(level)
                if id >= len(self.dir_parent):
                    self.dir_parent.append(None)
                self.dir_parent[id] = pdf_writer.addBookmark(level+' '+title, page-1, self.dir_parent[id-1])
                
                print(" * [%d/%d finished] level: %s(%d), title: %s, page: %d" % (i+1, num_all_lines, level, id, title, page))
            
            if self.out_path is None:
                self.out_path = self.pdf_path[:-4] + '(书签).pdf'
            with open(self.out_path, 'wb') as out_pdf:
                pdf_writer.write(out_pdf)
                print("---------------------------------------------------------------------------")
                print(" * Save: %s" % self.out_path)
                print("---------------------------------- Done! ----------------------------------")

if __name__ == '__main__':
    
    input_num = len(sys.argv)
    assert(input_num > 3)
    
    if input_num > 4:
        opath = sys.argv[4]
    else:
        opath = None

    if input_num > 5:
        mark = sys.argv[5]
    else:
        mark='.'

    pdg = PdfDirGenerator(
        pdf_path=sys.argv[1],
        txt_path=sys.argv[2],
        offset=int(sys.argv[3]), # 一般是目录结束页的页数
        out_path=opath,
        levelmark=mark
    )

    pdg.run()
